from database.mapper import stock_day_k
from utils import parse_util, caiptal_util, date_util, cache_util, statisticalutil, dict_util


def train(code: str) -> dict:
    key = 'train' + code + 'module'
    itv_train_data = cache_util.get_doc_expire(key)
    if itv_train_data is None:
        itv_train_data = {}
        # 一、训练百日资金流动模型
        main_100 = caiptal_util.flow_hundred(code)
        main_100.reverse()
        main_beg_rq = main_100[0].get('rq')
        main_end_rq = main_100[len(main_100) - 1].get('rq')
        sdk_interval = stock_day_k.select_interval(code, main_beg_rq, main_end_rq)
        sdk_interval.reverse()
        # 数据日期校对
        sdk_beg_rq = date_util.parse_date_str(sdk_interval[0].get('rq'))
        sdk_end_rq = date_util.parse_date_str(sdk_interval[len(sdk_interval) - 1].get('rq'))
        if sdk_beg_rq != main_beg_rq or sdk_end_rq != main_end_rq:
            print('train rq error')
            return {}
        if len(main_100) != len(sdk_interval):
            print('data len error')
            return {}
        for buy_interval in range(2, 15):
            for sell_interval in range(2, 15):
                amt_change = []  # 资金变化
                buy_roll_main_io = []  # 买入滚动主力净流入
                sell_roll_main_io = []  # 买入滚动主力净流入
                buy_rq_list = []  # 购买日期
                sel_rq_list = []  # 卖出日期
                buy_price_list = []  # 购买价格
                sel_price_list = []  # 卖出价格
                buy_itv_sum = 0  # 购买区间累计
                sell_itv_sum = 0  # 卖出区间累计
                amount = 100000
                hand = 0
                for i in range(len(main_100)):
                    buy_itv_sum += main_100[i].get('main_io')
                    sell_itv_sum += main_100[i].get('main_io')
                    if i >= buy_interval:
                        buy_itv_sum -= main_100[i - buy_interval].get('main_io')
                    if i >= sell_interval:
                        sell_itv_sum -= main_100[i - sell_interval].get('main_io')
                    buy_roll_main_io.append(round(buy_itv_sum / 100000000, 2))
                    sell_roll_main_io.append(round(sell_itv_sum / 100000000, 2))
                    if buy_itv_sum > 0 and sell_itv_sum > 0 and hand == 0 and i >= buy_interval:  # 买入
                        buy_rq_list.append(str(sdk_interval[i].get('rq')))  # 买入日期
                        sp = sdk_interval[i].get('sp')
                        buy_price_list.append(sp)
                        hand = round(amount / sp / 100, 0)
                        amount -= hand * sp * 100
                        # print(str(buy_interval) + '买:' + str(sdk_interval[i].get('rq')) + ':' + str(sp) + ':' + str(
                        #     buy_itv_sum))
                    elif sell_itv_sum < 0 and hand != 0 and i >= sell_interval:
                        sp = sdk_interval[i].get('sp')
                        amount += round(hand * sp * 100, 2)
                        hand = 0  # 手数清零
                        amt_change.append(amount)
                        sel_rq_list.append(str(sdk_interval[i].get('rq')))
                        sel_price_list.append(sp)
                        # print(str(sell_interval) + '卖:' + str(sdk_interval[i].get('rq')) + ':' + str(sp) + ':' + str(
                        #     sell_itv_sum))
                if hand != 0:  # 最新一天卖掉
                    amt_change.append(round(hand * sdk_interval[len(sdk_interval) - 1]['sp'] * 100, 2))
                    sel_rq_list.append(str(sdk_interval[len(sdk_interval) - 1].get('rq')))
                    sel_price_list.append(sdk_interval[len(sdk_interval) - 1].get('sp'))
                itv_train_data[str(buy_interval) + ':' + str(sell_interval)] = {
                    'buy_roll_main_io': buy_roll_main_io,
                    'sell_roll_main_io': sell_roll_main_io,
                    'amt_change': amt_change,
                    'buy_rq_list': buy_rq_list,
                    'buy_price_list': buy_price_list,
                    'sel_price_list': sel_price_list,
                    'sel_rq_list': sel_rq_list}
        # 计算所需数据
        champion_max = 0  # 总冠军
        champion_min = 999999999
        record_max = 0  # 最高纪录
        record_min = 999999999
        champion_max_inv = 0
        champion_min_itv = 0
        record_max_itv = 0
        record_min_itv = 0
        champion_bs_max = 0  # buy 和 sell相同的交易日 收益最大
        # 盈利排序
        amt_sum_sort = []
        for itv in itv_train_data.keys():
            itv_info = itv_train_data.get(itv)
            amt_change = itv_info.get('amt_change')
            if len(amt_change) == 0:
                continue
            if champion_max < amt_change[len(amt_change) - 1]:
                champion_max = amt_change[len(amt_change) - 1]
                champion_max_inv = itv
            if champion_min > amt_change[len(amt_change) - 1]:
                champion_min = amt_change[len(amt_change) - 1]
                champion_min_itv = itv
            amt_sum_sort.append({'key': itv, 'amt': amt_change[len(amt_change) - 1]})
            for ac in amt_change:
                if record_max < ac:
                    record_max = ac
                    record_max_itv = itv
                if record_min > ac:
                    record_min = ac
                    record_min_itv = itv
            # 二次加工折线图数据
            buy_roll_main_io = itv_info.get('buy_roll_main_io')
            sell_roll_main_io = itv_info.get('sell_roll_main_io')
            # 获取缩放比例
            main_io_max = max(max(buy_roll_main_io), max(sell_roll_main_io))
            main_io_min = min(min(buy_roll_main_io), min(sell_roll_main_io))
            sp_max = max(parse_util.list_dict_to_list(sdk_interval, 'sp'))
            sp_min = min(parse_util.list_dict_to_list(sdk_interval, 'sp'))
            bs = round((main_io_max - main_io_min) / (sp_max - sp_min), 8)
            x0 = []
            x1 = []
            x2 = []
            x3 = []
            y0 = []
            y1 = []
            y2 = []
            y3 = []
            v0 = []
            v1 = []
            v2 = []
            v3 = []
            for num in range(len(buy_roll_main_io)):
                x0.append(num)
                x1.append(num)
                x2.append(num)
                x3.append(num)
                y0.append(sp_min)
                y1.append(sdk_interval[num].get('sp'))
                y2.append(round(buy_roll_main_io[num] / bs, 2) + sp_min)
                y3.append(round(sell_roll_main_io[num] / bs, 2) + sp_min)
                if (num + 1) % 5 == 0:
                    v0.append(date_util.parse_date_str(sdk_interval[num].get('rq'))[-5:])
                    v1.append(sdk_interval[num].get('sp'))
                    v2.append(str(buy_roll_main_io[num]) + '亿')
                    v3.append(str(sell_roll_main_io[num]) + '亿')
                else:
                    v0.append('')
                    v1.append('')
                    v2.append('')
                    v3.append('')
            itv_train_data.get(itv)['broken_line_x'] = [x0, x1, x2, x3]
            itv_train_data.get(itv)['broken_line_y'] = [y0, y1, y2, y3]
            itv_train_data.get(itv)['broken_line_v'] = [v0, v1, v2, v3]
        # 计算最佳匹配模型
        amt_sum_sort = sorted(amt_sum_sort, key=lambda x: x.get('amt'), reverse=True)
        # 价格涨幅
        price_zf = round(
            (sdk_interval[len(sdk_interval) - 1].get('sp') - sdk_interval[0].get('sp')) / sdk_interval[0].get('sp'), 4)
        buy_x = []
        sel_x = []
        amt_compare = 0
        amt_flag = 0
        for i in range(min(50, len(amt_sum_sort))):
            if (price_zf <= 0 and (amt_sum_sort[i].get('amt') > 100000)) \
                    or round((amt_sum_sort[i].get('amt') - 100000) / 100000, 2) / price_zf > 0.7:
                amt_sum_sort[i]['zf'] = round((amt_sum_sort[i].get('amt') - 100000) / 100000, 2)
                amt_sum_sort[i]['buy'] = int(amt_sum_sort[i].get('key').split(':')[0])
                amt_sum_sort[i]['sell'] = int(amt_sum_sort[i].get('key').split(':')[1])
            if i == 0:
                amt_compare = round(amt_sum_sort[i].get('amt') * 0.8, 2)
            if i > 0 and (amt_sum_sort[i].get('amt') < amt_compare or
                          ((amt_sum_sort[i].get('amt') - 100000) * 1.08 + 100000) < amt_sum_sort[i - 1].get('amt')):
                # print(amt_sum_sort[i])
                amt_flag = 1
            # 计算收益离散性能
            if i <= 15 and amt_flag == 0:
                buy_x.append(int(amt_sum_sort[i].get('key').split(':')[0]))
                sel_x.append(int(amt_sum_sort[i].get('key').split(':')[1]))
            amt_sum_sort[i]['rank'] = i + 1
        buy_ls = round(statisticalutil.discrete_value(buy_x), 2)
        sel_ls = round(statisticalutil.discrete_value(sel_x), 2)
        for i in range(6):
            if buy_ls > 0.3:
                if i % 2 == 0:
                    buy_x.pop(buy_x.index(max(buy_x)))
                else:
                    buy_x.pop(buy_x.index(min(buy_x)))
                buy_ls = round(statisticalutil.discrete_value(buy_x), 2)
            if sel_ls > 0.3:
                if i % 2 == 0:
                    sel_x.pop(sel_x.index(max(sel_x)))
                else:
                    sel_x.pop(sel_x.index(min(sel_x)))
                sel_ls = round(statisticalutil.discrete_value(sel_x), 2)
        bs_arr = [[], []]
        if buy_ls <= 0.6:
            buy_x_dict = {}
            for i in buy_x:
                if i not in buy_x_dict:
                    buy_x_dict[i] = 1
                else:
                    buy_x_dict[i] = buy_x_dict[i] + 1
            buy_jfg = max(buy_x_dict, key=buy_x_dict.get)
            bs_arr[0].append(buy_jfg)
            if len(buy_x_dict) > 2:
                buy_x_dict.pop(buy_jfg)
                bs_arr[0].append(max(buy_x_dict, key=buy_x_dict.get))
        else:
            bs_arr[0].append(amt_sum_sort[0]['buy'])
        if sel_ls <= 0.6:
            sell_x_dict = {}
            for i in sel_x:
                if i not in sell_x_dict:
                    sell_x_dict[i] = 1
                else:
                    sell_x_dict[i] = sell_x_dict[i] + 1
            sel_jfg = max(sell_x_dict, key=sell_x_dict.get)
            bs_arr[1].append(sel_jfg)
            if len(sell_x_dict) > 2:
                sell_x_dict.pop(sel_jfg)
                bs_arr[1].append(max(sell_x_dict, key=sell_x_dict.get))
        else:
            bs_arr[1].append(amt_sum_sort[0]['sell'])
        amt_sum_sort_dict = parse_util.list_to_dict(amt_sum_sort, 'key')
        rank = 9999
        buy_jfg = 0
        sel_jfg = 0
        for b in bs_arr[0]:
            for s in bs_arr[1]:
                bsr = dict_util.get_default(amt_sum_sort_dict.get(str(b) + ':' + str(s)), 'rank', 10000)
                if bsr < rank:
                    rank = bsr
                    buy_jfg = b
                    sel_jfg = s

        itv_train_data['buy_ls'] = buy_ls
        itv_train_data['sel_ls'] = sel_ls
        itv_train_data['buy_jfg'] = buy_jfg
        itv_train_data['sel_jfg'] = sel_jfg

        flow_sum_dict = {}
        flow_sum = 0
        for i in range(15):
            flow_sum += main_100[len(main_100) - 1 - i].get('main_io')
            flow_sum_dict[i + 1] = flow_sum

        itv_train_data['beg_rq'] = sdk_beg_rq
        itv_train_data['end_rq'] = sdk_end_rq
        itv_train_data['champion_max'] = champion_max
        itv_train_data['champion_max_inv'] = champion_max_inv
        itv_train_data['champion_min'] = champion_min
        itv_train_data['champion_min_itv'] = champion_min_itv
        itv_train_data['record_max'] = record_max
        itv_train_data['record_max_itv'] = record_max_itv
        itv_train_data['record_min'] = record_min
        itv_train_data['record_min_itv'] = record_min_itv
        itv_train_data['beg_price'] = sdk_interval[0].get('sp')
        itv_train_data['end_price'] = sdk_interval[len(sdk_interval) - 1].get('sp')
        itv_train_data['main_100'] = main_100
        itv_train_data['flow_sum_dict'] = flow_sum_dict
        itv_train_data['amt_sum_sort'] = sorted(amt_sum_sort, key=lambda x: x.get('amt'), reverse=True)
        cache_util.set_doc_expire(key, str(itv_train_data), 1000)
    return itv_train_data


def del_module_doc_cache(code: str):
    cache_util.del_doc_cache('train' + code + 'module')


if __name__ == '__main__':
    # train("600089")
    # train("600036")
    # train("002389")  # 航天彩虹
    # train("600316")  # 洪都航空
    # train("600036")  # 招商银行
    # train("002466")  # 天齐锂业
    # train("002709")  # 天赐材料
    # train("300750")  # 宁德时代
    # train("600111")  # 北方稀土
    # train("000063")  # 中兴通讯
    # train("601318")  # 中国平安
    train("601728")  # 中国电信
    # train("601995")  # 中金公司
    # train("600089")  # 特变电工
    # print(train("600019"))  # 宝钢股份
    # print(train("002013"))  # 中航机电
    print("end")
